import logging
from concurrent.futures import ThreadPoolExecutor

from db import DB
from spider import Spider


class Main:
    def __init__(self):
        self.db = DB()
        self.thread_pool = ThreadPoolExecutor(max_workers=8)

    def thread_done(self, spyder: Spider, exception):
        # 线程回调函数
        if exception:
            # 抢课失败 重新添加抢课线程
            self.thread_pool.submit(spyder.start, self.thread_done)
        else:
            # 抢课成功写入数据库
            self.db.update_user_grab_lesson(spyder.uid, spyder.target_lesson_name)

    def cleanup(self):
        # 清理操作
        self.db.close()

    def start(self):
        cursor = self.db.get_failure_user_cursor()
        for user_info in cursor:
            spyder = Spider(user_info[0], user_info[1], user_info[2])
            self.thread_pool.submit(spyder.start, self.thread_done)


if __name__ == '__main__':
    logging.warning("回车三次终止选课")
    main = Main()
    main.start()
    input()
    input()
    input()
    main.cleanup()
